package iscas;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.Pipeline;

import java.util.HashMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BiFunction;
import java.util.function.Function;

public class LocalMapState<K, V> extends RedisLocalState {
    ConcurrentHashMap<K, V> local = new ConcurrentHashMap<>();
    BiFunction<V, V, V> merger;
    Function<String, K> keyParser;
    Function<String, V> valueParser;

    public LocalMapState(String name, BiFunction<V, V, V> merger, Function<String, K> keyParser,
                         Function<String, V> valueParser, long period) {
        super(name, period);
        this.merger = merger;
        this.keyParser = keyParser;
        this.valueParser = valueParser;
    }

    public void set(K k, V v) {
        local.put(k, v);
    }

    public V get(K k) {
        return local.get(k);
    }

    public void set(K k, Function<V, V> transform) {
        local.compute(k, (key, value) -> transform.apply(value));
    }

    public void set(K k, V defaultValue, Function<V, V> transform) {
        local.compute(k, (key, value) -> transform.apply(value == null ? defaultValue : value));
    }

    @Override
    boolean isEmpty() {
        return local.isEmpty();
    }

    @Override
    public void merge(Jedis jedis) {
        lock(jedis);
        HashMap<K, V> remote = new HashMap<>();
        jedis.hgetAll(name).forEach((k, v) -> remote.put(keyParser.apply(k), valueParser.apply(v)));
        Pipeline pipeline = jedis.pipelined();
        local.keySet().forEach(k -> {
            V v = local.remove(k);
            if (remote.containsKey(k)) {
                pipeline.hset(name, k.toString(), merger.apply(v, remote.get(k)).toString());
            } else {
                pipeline.hset(name, k.toString(), v.toString());
            }
        });
        pipeline.sync();
        unlock(jedis);
    }
}
